# write by hashaki
# first edit on 2018/11/13
# last change on 2019/04/24
# 常鸿量化机交易引擎

import shelve
import os
import sys
from collections import OrderedDict
from copy import copy
import traceback
import socket
import json

from datetime import datetime
class MainEngine:
    def __init__(self):
        self.socket=socket.socket()                 # socket
        self.socket_send_big=socket.socket()        # 发送数据到大策略的socket
        self.tick=None                              # 最新的tick {'symbol1':tick1,'symbol2':tick2,...}
        self.time=None                              # 最近的时间
        self.depths=None                            # 最新的深度数据
        self.bars=None                              # 输入的K线数据

        self.size=None                              # 合约大小
        self.slip=None                              # 滑点
        self.rate=None                              # 佣金比例
        self.cash=None                              # 现金流
        self.pos=0                                  # 持仓
        self.freeze=None                            # 冻结
        self.capital=None                           # 总资产

        self.strategyMat=None                       # 策略矩阵
        self.accountCapital=None                    # 账户信息，现金，持仓，冻结

        self.resultList=None                        # 策略传出来的信号结果

        self.port=None                              # 推送数据的端口
        self.address=None                           # 推送数据的地址

        self.__activate=None                        # 是否开始运行交易引擎

        self.orderDict={}                           # 下单字典
        self.isSave=False                           # 是否要保存每个小策略的下单结果
        self.shelvePath=os.getcwd()+'/data/shelve'  # 小策略持久化数据保存地址

        self.capitalList=[]                         # 总资产列表
        self.posList=[]                             # 持仓列表

        self.bars_to_strategy_time=0                # 用于策略获取到一次历史行情数据的计数器

    # ----------------------------------------------------------------------------
    # 数据输入模块
    # 获取最新的tick
    def tickIn(self,tick):
        '''tick数据输入tick={'symbol1':tick1,'symbol2':tick2,...}'''
        self.tick=tick
    
    def depthIn(self,depths):
        '''depth数据输入depth={'symbol1':{'asks':[[price,volume]],'bids':[[price,volume]],'symbol2':...} '''
        self.depths=depths

    def barsIn(self,bars):
        '''K线数据输入，{'symbol1':{'time'；datetime,'volume':float,'frequence':str，'open':float,'high':float,'low':float,'close':float}}'''
        self.bars=bars
    
    def tradeSetting(self,size,slip,rate):
        '''交易引擎设置'''
        self.size=size                  # 合约大小
        self.slip=slip                  # 滑点
        self.rate=rate                  # 佣金比例
    
    def strategyIn(self,strategyMat):
        '''策略矩阵{'strategy1':['tick',stategy1],'strategy2':['depth',strategy2],...} '''
        self.strategyMat=strategyMat
    
    def accountCapitalIn(self,msg):
        '''账户信息输入，包括现金，持仓，冻结,msg={'cash':float,'pos':{'symbol1':float,'symbols2':float,...},'freeze':{'symbol1':float,'symbol2':float,...} '''
        self.cash=msg['cash']
        self.pos=msg['pos']
        self.freeze=msg['freeze']
    
    def accountSimulate(self,cash,pos):
        '''模拟交易时候设置的现金和持仓'''
        self.cash=cash
        self.pos=pos
        self.capital=cash

    def engingSetting(self,address,port):
        '''交易引擎的设置,address喝port为获取数据推送的端口地址'''
        self.port=port
        if not address:
            self.address='127.0.0.1'
        else:
            self.address=address

    # -----------------------------------------------------------------------------
    # 数据处理
    def getStrategyOrder(self,strategy):
        '''获取策略文件返回的下单数据'''
        self.orderDict=strategy.returnOrder()
        if self.orderDict:
            self.orderDict['time']=datetime.now()
        else:
            pass

    def calculateCapital_simulate(self):
        '''资产计算,计算每一笔成交后持仓和现金和总资产'''
        if not self.orderDict:
            return

        side=self.orderDict['side']
        price=self.orderDict['price']
        volume=self.orderDict['volume']

        if side=='buy':
            if self.cash<10:
                print('持有现金不够下单')
                return

            if self.cash<price*volume:
                # 梭哈的情况
                self.pos+=self.cash/price
                self.cash-=self.pos * price
                self.capital=self.cash+price*self.pos
            
            else:
                self.pos+=volume
                self.cash-=price*volume
                self.capital=self.cash+price*self.pos
        
        else:
            if self.pos==0 or self.pos<0.01:
                print('持有仓位太小不足以卖出，当前持仓为：',self.pos)
                return

            elif self.pos<volume:
                # 清仓的情况
                self.cash+=self.pos*price
                self.pos=0
                self.capital=self.cash
            
            else:
                self.cash+=volume*price
                self.pos-=volume
                self.capital=self.cash+self.pos*price
        
        # 更新总资产列表和持仓列表
        self.capitalList.append(self.capital)
        self.posList.append(self.pos)

    def shelveTheOrder(self):
        '''持久化下单数据'''
        if self.isSave:
            if os.path.exists(self.shelvePath):
                # 持久化目录存在的情况
                path=self.shelvePath+'/order_base'
                order_base=shelve.open(path,writeback=True)
                if len(order_base)==0:
                    temp=[]
                    order_base['order']=temp
                    order_base['order'].append(self.orderDict)
                    order_base.close()
                    self.orderDict.clear()
                else:
                    order_base['order'].append(self.orderDict)
                    order_base.close()
                    self.orderDict.clear()
            else:
                data_path=os.getcwd()+'/data'
                if os.path.exists(data_path):
                    # 有data目录但没有持久化目录的情况
                    os.mkdir(self.shelvePath)
                    path=self.shelvePath+'/order_base'
                    order_base=shelve.open(path,writeback=True)
                    if len(order_base)==0:
                        temp=[]
                        order_base['order']=temp
                        order_base['order'].append(self.orderDict)
                        order_base.close()
                        self.orderDict.clear()
                    else:
                        order_base['order'].append(self.orderDict)
                        order_base.close()
                        self.orderDict.clear()
                else:
                    # 连data 目录都没有的情况
                    os.mkdir(data_path)
                    os.mkdir(self.shelvePath)
                    path=self.shelvePath+'/order_base'
                    order_base=shelve.open(path,writeback=True)
                    if len(order_base)==0:
                        temp=[]
                        order_base['order']=temp
                        order_base['order'].append(self.orderDict)
                        order_base.close()
                        self.orderDict.clear()
                    else:
                        order_base['order'].append(self.orderDict)
                        order_base.close()
                        self.orderDict.clear()
                    
        else:
            self.orderDict.clear()
            
    # -----------------------------------------------------------------------------
    # 跑起来哦
    def start(self):
        '''启动交易引擎'''
        self.resultList=[]
        self.__activate=True
    
    def stop(self):
        '''关闭交易引擎'''
        self.__activate=False

    def run(self):
        '''总运行程序'''
        self.socket.connect((self.address,self.port))            # 连接数据端口
        self.socket_send_big.bind((self.address,self.port+1))    # 创建大策略连接
        self.socket_send_big.listen(1)
        client,addr=self.socket_send_big.accept()

        while self.__activate:
            data_str=self.socket.recv(500000).decode()           # 由于要加载过去24小时的一分钟K线，所以这里的缓存要开很大
            
            try:
                data=json.loads(data_str)
            except json.decoder.JSONDecodeError:
                continue

            if data['type']=='kline':
                print('成功加载24小时历史数据')
                self.bars=data['data']                           # self.bars={'symbol':[{'time':timestamp,'open':,'close':,'high':,'low':,'volume':},...]},K线数据只获取一次

            for strategy in self.strategyMat:
                temp=list(strategy.values())[0]

                if self.bars_to_strategy_time==0:
                    temp[1].history_bars=self.bars                # 初始化历史行情数据到策略

                if data['type']==temp[0]:
                    temp[1].pos=self.pos                          # 更新持仓信息
                    temp[1].cash=self.cash                        # 更新持有现金信息
                    temp[1].capital=self.capital                  # 更新总资产
                    temp[1].next(data['data'])                    # temp[1]就是策略文件
                    self.getStrategyOrder(temp[1])
                    self.calculateCapital_simulate()

                    if not self.orderDict:
                        # 策略没有输出下单指令，直接跳过持久化过程
                        continue
                    else:
                        for name in strategy:
                            self.orderDict['strategy']=name
                        temp=str(self.orderDict)
                        client.sendall(temp.encode('utf8'))
                        self.shelveTheOrder()

                else:
                    continue
                
            self.bars_to_strategy_time+=1
    
    def runSmallStrategy(self):
        '''没有加进大策略的运行'''
        self.socket.connect((self.address,self.port))            # 连接数据端口

        while self.__activate:
            data_str=self.socket.recv(500000).decode()           # 由于要加载过去24小时的一分钟K线，所以这里的缓存要开很大
            
            try:
                data=json.loads(data_str)
            except json.decoder.JSONDecodeError:
                continue

            if data['type']=='kline':
                print('成功加载24小时历史数据')
                self.bars=data['data']                           # self.bars={'symbol':[{'time':timestamp,'open':,'close':,'high':,'low':,'volume':},...]},K线数据只获取一次

            for strategy in self.strategyMat:
                temp=list(strategy.values())[0]

                if self.bars_to_strategy_time==0:
                    temp[1].history_bars=self.bars                # 初始化历史行情数据到策略

                if data['type']==temp[0]:
                    temp[1].pos=self.pos                          # 更新持仓信息
                    temp[1].cash=self.cash                        # 更新持有现金信息
                    temp[1].capital=self.capital                  # 更新总资产
                    temp[1].next(data['data'])                    # temp[1]就是策略文件
                    self.getStrategyOrder(temp[1])
                    self.calculateCapital_simulate()

                    if not self.orderDict:
                        # 策略没有输出下单指令，直接跳过持久化过程
                        continue
                    else:
                        for name in strategy:
                            self.orderDict['strategy']=name
                        temp=str(self.orderDict)
                        self.shelveTheOrder()

                else:
                    continue
                
            self.bars_to_strategy_time+=1
                